# # Reference data generation of xc_deriv

# +
from pyscf import gto, scf, dft, lib
import numpy as np

np.set_printoptions(6, suppress=True, linewidth=150)
# -

# ## Per-file Configuration

xctype = "MGGA"
xc = "TPSS"

# ## Molecule Definition

mol = gto.Mole(atom="O; H 1 0.94; H 1 0.94 2 104.5", spin=2, charge=2, basis="STO-3G").build()

ni = dft.numint.NumInt()

mf = dft.UKS(mol, xc=xc)
mf.grids.level = 0
mf.run()

ngrids = mf.grids.weights.size
ngrids

dm = mf.make_rdm1()
grids = mf.grids
ao = ni.eval_ao(mol, grids.coords, deriv=2)
rho = np.array([
    ni.eval_rho(mol, ao, dm[0], xctype=xctype),
    ni.eval_rho(mol, ao, dm[1], xctype=xctype),
])

rho.shape

np.save("rho", rho)

# ## Effective kernels

xc0 = dft.libxc._eval_xc(xc, rho, spin=1, deriv=3)
np.save("xc0", xc0)

exc, vxc, fxc, kxc = ni.eval_xc(xc, rho, deriv=3, spin=1)

np.save("exc_0", exc)
[np.save(f"vxc_{i}", g) for i, g in enumerate(vxc) if g is not None]
[np.save(f"fxc_{i}", g) for i, g in enumerate(fxc) if g is not None]
[np.save(f"kxc_{i}", g) for i, g in enumerate(kxc) if g is not None]

exc_eff, vxc_eff, fxc_eff, kxc_eff = ni.eval_xc_eff(xc, rho, deriv=3)

np.save("exc_eff", exc_eff)
np.save("vxc_eff", vxc_eff)
np.save("fxc_eff", fxc_eff)
np.save("kxc_eff", kxc_eff)
